package com.abyss.asynio;


import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.java.tuple.Tuple1;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.AsyncDataStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.async.ResultFuture;
import org.apache.flink.streaming.api.functions.async.RichAsyncFunction;


import java.sql.*;
import java.util.Arrays;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;
import java.util.function.Supplier;

/**
 * socket传入id
 * 异步去mysql查询id代表的信息
 */
public class AsyncIODemo {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        SingleOutputStreamOperator<Tuple1<Integer>> idTuple = env.socketTextStream("node1", 9999).map(new MapFunction<String, Tuple1<Integer>>() {
            @Override
            public Tuple1<Integer> map(String value) throws Exception {
                return Tuple1.of(Integer.parseInt(value));
            }
        });

        // 使用异步的形式去查询
        // 发起异步IO
        SingleOutputStreamOperator<MysqlResult> unorderedWaitDataStream = AsyncDataStream
                .unorderedWait(idTuple, new MyAsyncIOImpl(), 60000L, TimeUnit.MILLISECONDS);
        idTuple.print(">>>input id: ");
        unorderedWaitDataStream.print();

        env.execute();
    }

    public static class MyAsyncIOImpl extends RichAsyncFunction<Tuple1<Integer>, MysqlResult> {
        private Connection connection = null;
        private PreparedStatement ps = null;

        @Override
        public void open(Configuration parameters) throws Exception {
            super.open(parameters);
            Class.forName("com.mysql.jdbc.Driver");
            String url = "jdbc:mysql://node1:3306/flinkdemo?useUnicode=true&characterEncoding=utf-8&useSSL=false";
            this.connection = DriverManager.getConnection(url, "root", "123456");
            this.ps = this.connection.prepareStatement("SELECT id, username, name FROM user WHERE id = ?");
        }

        @Override
        public void close() throws Exception {
            super.close();
            if (this.ps != null) this.ps.close();
            if (this.connection != null) this.connection.close();
        }

        @Override
        public void asyncInvoke(Tuple1<Integer> input, ResultFuture<MysqlResult> resultFuture) {
            CompletableFuture.supplyAsync(new Supplier<MysqlResult>() {
                @Override
                public MysqlResult get() {
                    try {
                        ps.setInt(1, input.f0);
                        ResultSet resultSet = ps.executeQuery();
                        if (resultSet.next()) {
                            TimeUnit.SECONDS.sleep(5);
                            return new MysqlResult(resultSet.getInt("id"), resultSet.getString("username"), resultSet.getString("name"));
                        }
                    } catch (SQLException | InterruptedException throwables) {
                        throwables.printStackTrace();
                    }

                    return null;
                }
            }).thenAccept( (MysqlResult mr) -> {
                resultFuture.complete(Collections.singleton(mr));
            } );
        }
    }

    public static class MysqlResult {
        private int id;
        private String username;
        private String name;

        @Override
        public String toString() {
            return "MysqlResult{" +
                    "id=" + id +
                    ", username='" + username + '\'' +
                    ", name='" + name + '\'' +
                    '}';
        }

        public int getId() {
            return id;
        }

        public void setId(int id) {
            this.id = id;
        }

        public String getUsername() {
            return username;
        }

        public void setUsername(String username) {
            this.username = username;
        }

        public String getName() {
            return name;
        }

        public void setName(String name) {
            this.name = name;
        }

        public MysqlResult() {
        }

        public MysqlResult(int id, String username, String name) {
            this.id = id;
            this.username = username;
            this.name = name;
        }
    }
}
